package com.huqing.icu.service.vectorstore;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.huqing.icu.model.MilvusFileCollection;
import com.huqing.icu.utils.JsonUtils;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DescribeCollectionReq;
import io.milvus.v2.service.collection.response.DescribeCollectionResp;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * @Description milvus向量库存储
 * @Author huqing
 * @Date 2025/8/5 15:19
 **/
@Service
public class MilvusEmbeddingServiceImpl implements IVectorStoreService {

    private static final Logger logger = LoggerFactory.getLogger(IVectorStoreService.class);

    private final String COLLECTION_NAME = "huqing_collection";

    @Autowired
    private MilvusClientV2 milvusClientV2;
    @Autowired
    private OpenAiEmbeddingModel openAiEmbeddingModel;

    /*@Autowired
    private VectorStore vectorStore;*/


    /**
     * 创建文件阅读器
     *
     * @param fileName
     * @param resource
     * @return
     */
    @Override
    public DocumentReader createDocumentReader(String fileName, Resource resource) {
        fileName = fileName.toLowerCase();

        DocumentReader documentReader = null;
        if (fileName.endsWith(".json")) {
            //json文件
            documentReader = new JsonReader(resource);
        } else if (fileName.endsWith(".txt")) {
            //txt文件
            documentReader = new TextReader(resource);
            TextReader textReader = new TextReader(resource);
        } else if (fileName.endsWith(".pdf") || fileName.endsWith(".doc") ||
                fileName.endsWith(".docx") || fileName.endsWith(".ppt")
                || fileName.endsWith(".pptx") || fileName.endsWith(".html")) {
            //pdf、doc、docx、ppt、pptx、html文件
            documentReader = new TikaDocumentReader(resource);
        } else {
            throw new IllegalArgumentException("不支持此文件格式");

        }
        return documentReader;
    }

    /**
     * 阅读文件
     *
     * @param resource
     * @param fileName
     * @return
     */
    @Override
    public List<Document> read(Resource resource, String fileName) {
        DocumentReader documentReader = createDocumentReader(fileName, resource);
        List<Document> documentList = documentReader.read();
        //文本拆分器，将文本拆分成较小的块
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> newDocumentList = splitter.split(documentList);

        //给每个文本设置元数据，包括文件名
        for (Document document : newDocumentList) {
            Map<String, Object> metadata = document.getMetadata();
            metadata.put("fileName", fileName);
        }
        return newDocumentList;
    }

    /**
     * 把文本存到向量库
     *
     * @param documentList
     * @return
     */
    @Override
    public Integer saveVectorStore(List<Document> documentList, String fileName) {


        //vectorStore.add(documentList);
        Gson gson = new Gson();
        List<JsonObject> collectionDataList = new ArrayList(documentList.size());
        for (Document document : documentList) {
            float[] embed = openAiEmbeddingModel.embed(document.getText());

            //保存到Milvis
            MilvusFileCollection milvusFileCollection = new MilvusFileCollection();
            //milvusFileCollection.setId(0L);
            milvusFileCollection.setUser_id(0L);
            milvusFileCollection.setFile_name(fileName);
            milvusFileCollection.setFile_id(null);
            milvusFileCollection.setContent(document.getText());
            milvusFileCollection.setContent_vector(embed);
            JsonObject jsonObject = gson.fromJson(JsonUtils.obj2String(milvusFileCollection), JsonObject.class);
            collectionDataList.add(jsonObject);
        }
        InsertReq insertReq = InsertReq.builder().collectionName(COLLECTION_NAME).data(collectionDataList).build();
        logger.info("save to milvus  insertReq = {}", insertReq);
        InsertResp insertResp = milvusClientV2.insert(insertReq);
        logger.info("save to milvus  insertResp = {}", insertResp);
        return 1;
    }

    /**
     * 向量库搜索，这是springAI的向量搜索
     *
     * @param content             搜索的内容
     * @param similarityThreshold 文档匹配的最低相似度阈值
     * @param topK                按分数排名，取前几个
     * @param fileNameList
     * @return
     */
    /*@Override
    public List<Document> similarityVectorData(String content, Double similarityThreshold, Integer topK, List<Object> fileNameList) {
        //元数据过滤器，用于元数据筛选
        FilterExpressionBuilder filterExpressionBuilder = new FilterExpressionBuilder();
        Filter.Expression expression = filterExpressionBuilder.in("fileName", fileNameList).build();

        //从向量库搜索，filterExpression是基于元数据的过滤条件
        SearchRequest searchRequest = SearchRequest.builder().query(content).topK(topK).similarityThreshold(similarityThreshold).filterExpression(expression).build();
        //List<Document> documentList = vectorStore.similaritySearch(searchRequest);
        List<Document> documentList = null;
        return documentList;
    }*/

    /**
     * 这是milvus的sdk提供的向量搜索
     *
     * @param content
     * @param similarityThreshold
     * @param topK
     * @param fileNameList
     * @return
     */
    @Override
    public List<Document> similarityVectorData(String fileName, String content, Double similarityThreshold, Integer topK, List<Object> fileNameList) {


        //先将查询文本向量化
        float[] queryEmbedding = openAiEmbeddingModel.embed(content);
        FloatVec queryVector = new FloatVec(queryEmbedding);
        SearchReq searchReq = SearchReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(Collections.singletonList(queryVector))
                .topK(topK)
                .filter("file_name == \"fileName\" ")
                //.outputFields(Arrays.asList("color", "likes"))
                .build();
        SearchResp searchResp = milvusClientV2.search(searchReq);

        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();

        return null;
    }

    /**
     * 创建集合
     */
    public void createCollection() {

        //创建一个schema
        CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema();

        //为这个schema添加字段
        //主键字段
        schema.addField(AddFieldReq.builder().fieldName("id").dataType(DataType.VarChar).isPrimaryKey(true).autoID(true).build());

        //用户ID
        schema.addField(AddFieldReq.builder().fieldName("user_id").dataType(DataType.Int32).build());

        //文件名称
        schema.addField(AddFieldReq.builder().fieldName("file_name").dataType(DataType.VarChar).maxLength(10000).build());

        //文件ID
        schema.addField(AddFieldReq.builder().fieldName("file_id").dataType(DataType.Int32).build());

        //内容
        schema.addField(AddFieldReq.builder().fieldName("content").dataType(DataType.VarChar).maxLength(10000).build());

        //向量内容：向量维度定义1536，跟阿里巴巴embedding向量服务返回的维度保持一致
        schema.addField(AddFieldReq.builder().fieldName("content_vector").dataType(DataType.FloatVector).dimension(1536).build());

        //为向量字段建立索引
        IndexParam indexParam = IndexParam.builder().fieldName("content_vector").metricType(IndexParam.MetricType.COSINE).build();

        //集合名称
        CreateCollectionReq createCollectionReq = CreateCollectionReq.builder().collectionName("huqing_collection").collectionSchema(schema).indexParams(Collections.singletonList(indexParam)).build();

        logger.info("milvus create collection start");
        //创建集合
        milvusClientV2.createCollection(createCollectionReq);
        logger.info("milvus create collection success");

        DescribeCollectionReq describeCollectionReq = DescribeCollectionReq.builder().collectionName("quick_setup").build();
        DescribeCollectionResp describeCollectionResp = milvusClientV2.describeCollection(describeCollectionReq);
        logger.info("milvus query collection, resp = {}", JsonUtils.obj2String(describeCollectionResp));
    }
}
